import random

import numpy as np
import pickle


w1, word_to_index, index_to_word, word_size = pickle.load(open('./datas/word2vec.pkl', 'rb'))


def word_vector(word):
    return w1[word_to_index[word]]


def find_words(word, num):
    now_word_vector = word_vector(word)
    words_dic = {}
    for i in range(int(word_size)):
        other_word_vector = w1[i]
        theta_sum = np.dot(now_word_vector, other_word_vector)
        theta_den = np.linalg.norm(now_word_vector) * np.linalg.norm(other_word_vector)
        theta = theta_sum / theta_den
        word = index_to_word[i]
        words_dic[word] = theta
    words_sorted = sorted(words_dic.items(), key=lambda kv: kv[1], reverse=True)
    for word, value in words_sorted[:num]:
        print(word, value)


for i in range(3):
    n = random.randint(0, word_size)
    word = index_to_word[n]
    print(word)
    print(find_words(word, 10))
